from pydantic import BaseModel
from typing import List, Optional, Literal, Dict, Any
from datetime import datetime

# 原始的StreamResponse格式
class StreamResponse(BaseModel):
    type: Literal["message", "error", "complete"]
    content: str
    message_order: Optional[int] = None
    token_count: Optional[int] = None
    model_name: Optional[str] = None
    processing_time_ms: Optional[int] = None

# 自定义格式1：更详细的步骤信息
class DetailedStreamResponse(BaseModel):
    type: Literal["message", "error", "complete", "step_start", "step_end"]
    step: Optional[str] = None  # element_extraction, ggb_generation, html_generation, html_optimization
    content: str
    progress: Optional[float] = None  # 0.0 - 1.0
    timestamp: Optional[datetime] = None
    message_order: Optional[int] = None
    metadata: Optional[Dict[str, Any]] = None

# 自定义格式2：分层结构
class HierarchicalStreamResponse(BaseModel):
    event_type: Literal["processing", "result", "error", "complete"]
    stage: Optional[str] = None
    data: Dict[str, Any]
    context: Optional[Dict[str, Any]] = None

# 自定义格式3：简化格式
class SimpleStreamResponse(BaseModel):
    status: Literal["processing", "done", "error"]
    message: str
    data: Optional[Dict[str, Any]] = None

# 自定义格式4：前端友好格式
class FrontendStreamResponse(BaseModel):
    id: Optional[str] = None  # 消息ID
    type: Literal["text", "html", "commands", "error", "complete"]
    payload: str
    step_name: Optional[str] = None
    step_progress: Optional[int] = None  # 1-4 (对应4个步骤)
    is_final: bool = False
    created_at: Optional[datetime] = None

# 自定义格式5：结构化数据格式
class StructuredStreamResponse(BaseModel):
    message_id: str
    conversation_id: int
    response_type: Literal["reasoning", "elements", "commands", "html", "optimized_html", "error", "complete"]
    content: str
    step_info: Optional[Dict[str, Any]] = None
    processing_stats: Optional[Dict[str, Any]] = None 